package com.bigdata.hive.hooks;

import org.apache.hadoop.hive.ql.hooks.ExecuteWithHookContext;
import org.apache.hadoop.hive.ql.hooks.HookContext;

/**
 * DESC
 *
 * @author Cyber
 * <p> Created By 2025/5/22
 * @version 1.0
 */

public class CustomExecuteHook implements ExecuteWithHookContext {

    /*
<property>
    <name>hive.exec.post.hooks</name>
    <value>com.bigdata.hive.hooks.CustomExecuteHook</value>
</property>
<property>
    <name>hive.exec.pre.hooks</name>
    <value>com.bigdata.hive.hooks.CustomExecuteHook</value>
</property>
<property>
    <name>hive.exec.failure.hooks</name>
    <value>com.bigdata.hive.hooks.CustomExecuteHook</value>
</property>

     */
    @Override
    public void run(HookContext hookContext) throws Exception {
        // 仅处理 HookType 阶段的钩子
        if (hookContext.getHookType() != HookContext.HookType.PRE_EXEC_HOOK
                && hookContext.getHookType() != HookContext.HookType.POST_EXEC_HOOK
                && hookContext.getHookType() != HookContext.HookType.ON_FAILURE_HOOK) {
            return;
        }

        // 获取 SQL 和用户信息
        String query = hookContext.getQueryPlan().getQueryString();
        String ID = hookContext.getQueryPlan().getQueryId();
        String user = hookContext.getUserName();
//        String user = SessionState.get().getUserName();
        // 捕获执行前日志
        if (hookContext.getHookType() == HookContext.HookType.PRE_EXEC_HOOK) {
            System.out.println("[PRE-HOOK]" + ID + " User: " + user + ", Query: " + query);
        }

        //  捕获执行错误日志
        if (hookContext.getHookType() == HookContext.HookType.ON_FAILURE_HOOK) {
            Throwable exception = hookContext.getException();
            System.out.println("[ON-FAILURE-HOOK]" + ID + " User: " + user + ", Query: " + query + ", 执行失败【Exception: " + exception + "】");
        }
        // 捕获执行后日志
        if (hookContext.getHookType() == HookContext.HookType.POST_EXEC_HOOK) {
            Throwable exception = hookContext.getException();
            int status = (exception == null) ? 0 : 1;
            System.out.println("[POST-HOOK]" + ID + " User: " + user + ", Query: " + query + "，执行成功【" + status + "】");
            if (exception != null) {
                System.err.println("Error: " + exception.getMessage());
            }
        }
    }
}
